#!/usr/bin/env python

# imports
import re
import sys

# create patterns to match number types
bin1pattern = re.compile('^[0,1]+b$')
bin2pattern = re.compile('^0b[0,1]+$')
bin3pattern = re.compile('^[-]+[0,1]+b$')
bin4pattern = re.compile('^[-]+0b[0,1]+$')
octpattern  = re.compile('^[-]?0[0-7]+$')
hexpattern  = re.compile('^0x[0-9,a-f,A-F]+$')
decpattern  = re.compile('^[-]?[0-9]+$')
charpattern = re.compile('^[a-z]$', re.I)


# usage
def urtPluginDescription(pyew, doprint=True):
  ''' Describes all functions provided by the URT plugin '''
  if doprint:
    print('[*] Ultimate Reversing Tool plugin provides:')
    print('    urt:    Prints this informational message')
    print('    base:   Prints a number in 4 bases                   eg: base 123')
    print('    oper:   Performs an operation on numbers or chars    eg: oper 101b xor 111b')
    print('            Binary: xor, and, or, mul, div,')
    print('                    add, sub, mod, exp, shl, shr')
    print('            Unary : not, neg')
    print('    caves:  Finds traditional code caves                 eg: caves <min_size>')
    print('    fstack: A fake stack to play with                    eg: fstack <32|64>')
  return
  

# convert string argument to an integer
def str2int(stringArg):
  # check number against patterns
  if re.match(bin1pattern, stringArg):
    numStr = stringArg[:-1]
    numInt = int(numStr, 2)
  elif re.match(bin2pattern, stringArg):
    numInt = int(stringArg, 2)
  elif re.match(bin3pattern, stringArg):
    numStr = stringArg[1:-1]
    numInt = int(numStr, 2)
  elif re.match(octpattern, stringArg):
    numInt = int(stringArg, 8)
  elif re.match(hexpattern, stringArg):
    numInt = int(stringArg, 16)
  elif re.match(decpattern, stringArg):
    numInt = int(stringArg, 10)
  elif re.match(charpattern, stringArg):
    numInt = ord(stringArg)
  else:
    print('[error] invalid number format')
    return False
  return numInt
  

# conversion function
def convertBase(pyew, args, doprint=True):
  ''' Prints a number or char in hex, decimal, octal and binary '''
  # check for args
  try:
    if len(args) != 0:
      number = args[0]
  except IndexError:
    if doprint:
      print('[error] requires an argument')
      return
    else:
      return('[error] requires an argument')

  # convert string to int
  baseInt = str2int(number)

  # make binary output nibble aligned
  if bin(baseInt)[0] == "-":
    binNoHeader = bin(baseInt)[3:]
    binNum = '-'
    remainder = len(binNoHeader)%4
    if remainder != 0:
      binNum += binNoHeader[:remainder]
      binNum += ' '
  else:
    binNoHeader = bin(baseInt)[2:]
    binNum = ''
    remainder = len(binNoHeader)%4
    if remainder != 0:
      binNum += binNoHeader[:remainder].zfill(4)
      binNum += ' '
  for i in range(remainder, len(binNoHeader), 4):
    binNum += binNoHeader[i:i+4]
    binNum += ' '

  result = '[base] hex: ' + hex(baseInt) + '  dec: ' + str(baseInt) + '  oct: ' + oct(baseInt) + '  bin: ' + binNum
  if doprint:
    print(result)
    return True
  else:
    return result


def operation(pyew, args, doprint=True):
  ''' Performs a logical or mathematical operation '''
  # valid operation tables
  binaryOperators = ('xor', 'or', 'and', 'mul', 'div', 'add', 'sub',
                     'mod', 'exp', 'shl', 'shr')
  unaryOperators = ('not', 'neg',)

  # check arg count
  try:
    # binary operation ?
    if len(args) == 3:
      operand1 = args[0]
      operator = args[1]
      operand2 = args[2]

      # convert string to int
      binOpInt1 = str2int(operand1)
      binOpInt2 = str2int(operand2)
      if binOpInt1 is False or binOpInt2 is False:
          return False

      # check operation is valid then execute
      if (operator in binaryOperators) is True:
        if operator == 'xor':
           convertBase(pyew, [str(binOpInt1 ^ binOpInt2),])
        if operator == 'or':
           convertBase(pyew, [str(binOpInt1 | binOpInt2),])
        if operator == 'and':
           convertBase(pyew, [str(binOpInt1 & binOpInt2),])
        if operator == 'mul':
           convertBase(pyew, [str(binOpInt1 * binOpInt2),])
        if operator == 'div':
           convertBase(pyew, [str(binOpInt1 / binOpInt2),])
        if operator == 'add':
           convertBase(pyew, [str(binOpInt1 + binOpInt2),])
        if operator == 'sub':
           convertBase(pyew, [str(binOpInt1 - binOpInt2),])
        if operator == 'mod':
           convertBase(pyew, [str(binOpInt1 % binOpInt2),])
        if operator == 'exp':
           convertBase(pyew, [str(binOpInt1 ** binOpInt2),])
        if operator == 'shl':
           convertBase(pyew, [str(binOpInt1 << binOpInt2),])
        if operator == 'shr':
           convertBase(pyew, [str(binOpInt1 >> binOpInt2),])
      else:
        # invalid operator
        if doprint:
          print('[error] invalid operator')
          return
        else:
          return('[error] invalid operator')

    # unary operation ?
    elif len(args) == 2:
      operator = args[0]
      operand1 = args[1]

      # convert string to int
      unOpInt1 = str2int(operand1)
      if unOpInt1 is False:
          return False

      # check operation is valid then execute
      if (operator in unaryOperators) is True:
        if operator == 'not':
           convertBase(pyew, [str(~ unOpInt1),])
        if operator == 'neg':
           convertBase(pyew, [str(- unOpInt1),])
      else:
        # invalid operator
        if doprint:
          print('[error] invalid operator: ', operator)
          return
        else:
          return('[error] invalid operator', operator)

  # invalid number of args
  except IndexError:
    if doprint:
      print('[error] requires 2 or 3 arguments')
      return
    else:
      return('[error] requires 2 or 3 arguments')


def findCodeCaves(pyew, args, doprint=True):
  ''' Locates traditional code caves '''
  # see if user provided a minimum cave size
  if len(args) == 1:
    minCaveSize = args[0]
  else:
    print('[error] one numeric argument required')
    return False

  # flag to test whether the section is executable
  execFlag = 0x20000000
  caveRegex = '\x00{' + minCaveSize + ',}'
  cavePattern = re.compile(caveRegex)
  for section in pyew.pe.sections:
    if (section.Characteristics & execFlag) != 0:
      executable = 'Yes'
    else:
      executable = 'No'
    # save location
    savedOffset = pyew.offset
    pyew.offset = section.PointerToRawData
    print('[section] Name: %s  Executable: %s' % (section.Name.rstrip('\x00'), executable))
    while (pyew.offset <= (section.PointerToRawData + section.SizeOfRawData - pyew.pe.OPTIONAL_HEADER.FileAlignment)):
      pyew.f.seek(pyew.offset)
      buf = pyew.f.read(pyew.pe.OPTIONAL_HEADER.FileAlignment)
      # search in buf
      caveIterator = cavePattern.finditer(buf)
      for match in caveIterator:
        print('[cave]    Offset: %s  Length: %d' % (hex(match.start() + pyew.offset), len(match.group())))
      pyew.offset += pyew.pe.OPTIONAL_HEADER.FileAlignment

    # reset to previous location
    pyew.offset = savedOffset
    pyew.seek(pyew.offset)


# fake stack object, instantiated by fstack()
class fakeStack:
  multiplier = 0
  base = ''
  top = ''
  def __init__(self, bits):
    if bits == '64':
      self.multiplier = 8
      self.base = 'rbp'
      self.top = 'rsp'
      self.stackStart = 0x770000000000
      self.savedBase = 0x780000000000
    elif bits == '32':
      self.multiplier = 4
      self.base = 'ebp'
      self.top = 'esp'
      self.stackStart = 0x77000000
      self.savedBase = 0x78000000
    else:
      print('[error] bits must be 32 or 64')
      exit()

    self.stack = [self.savedBase,]
    self.stackIndex = 0
    self.basePointer = self.stackStart
    self.stackPointer = self.stackStart
    self.printStack()


  def getStackPointer(self):
    return self.stack[self.stackPointer]


  def push(self, a):
    ''' push a value onto the stack '''
    #if
    self.stack.append(a)
    self.stackIndex += 1
    self.stackPointer -= self.multiplier
    self.printStack()
    return True

  def pop(self):
    ''' pop a value off the stack '''
    self.stack.pop()
    self.stackIndex -= 1
    self.stackPointer += self.multiplier
    self.printStack()
    return True

  def printStack(self):
    addrFormat = '0x%0.' + str(self.multiplier * 2) + 'x'
    for i in range(0, self.stackIndex +1 ):
      if i == 0:
        if self.stackPointer == self.basePointer:
          print('[%s][%s]   [' + addrFormat + ']') % (self.base, self.top, self.stack[i])
        else:
          print('[%s]        [' + addrFormat + ']') % (self.base, self.stack[i])
      elif i == self.stackIndex:
        print('     [%s]   [' + addrFormat + ']') % (self.top, self.stack[i])
      else:
        print('             [' + addrFormat + ']') % (self.stack[i],)
      i += 1
    print ''
    print('%s:' + addrFormat + ' %s:' + addrFormat) % (self.base, self.basePointer, self.top, self.stackPointer)


def fstack(pyew, args):
  ''' A fake stack you can play with '''
  if len(args) == 1:
    print('* try \'push <num>\' or \'pop\'')
    fs = fakeStack(args[0])
  else:
    print('[error] only one argument required to initialise fake stack')
    return
  while True:
    line = raw_input('fakeStack> ')
    cmds = line.split(' ')
    if len(cmds) == 2 and cmds[0] == 'push':
      val = str2int(cmds[1])
      if val is not False:
        fs.push(val)
    if len(cmds) == 1 and cmds[0] == 'pop':
      fs.pop()
    if len(cmds) == 1 and (cmds[0] == 'quit' or cmds[0] == 'q'):
      break


# end
functions = {"urt":urtPluginDescription,
             "base":convertBase,
             "oper":operation,
             "caves":findCodeCaves,
             "fstack":fstack}


